package org.example.aop;

import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.example.annotation.RateLimit;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ResponseStatusException;

import java.util.concurrent.ConcurrentHashMap;

@Component
@Aspect
@Slf4j
public class RateLimitAspect {

    // 使用线程安全的 ConcurrentHashMap 缓存 RateLimiter 实例
    private final ConcurrentHashMap<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();

    @Around("@annotation(rateLimit)")
    public Object around(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        log.info("进入限流切面，方法名: {}", joinPoint.getSignature().toLongString());
        // 获取注解参数
        double permitsPerSecond = rateLimit.permitsPerSecond();
        String methodName = joinPoint.getSignature().toLongString();

        // 获取或创建 RateLimiter 实例
        RateLimiter rateLimiter = rateLimiterMap.computeIfAbsent(
                methodName,
                key -> RateLimiter.create(permitsPerSecond)
        );

        // 尝试获取令牌
        if (rateLimiter.tryAcquire()) {
            log.info("允许方法执行");
            // 允许执行方法
            return joinPoint.proceed();
        } else {
            log.info("触发限流，返回错误信息");
            // 触发限流，返回错误信息
            throw new ResponseStatusException(HttpStatus.TOO_MANY_REQUESTS, rateLimit.message()
            );
        }
    }


}
